import torch
from collections import OrderedDict

from basicsr.archs import build_network
from basicsr.losses import build_loss
from basicsr.utils import get_root_logger
from basicsr.utils.registry import MODEL_REGISTRY
from .sr_mix_model import SRMixModel
from copy import deepcopy


@MODEL_REGISTRY.register()
class SRGANRefModel(SRMixModel):
    """SRGAN model for single image super-resolution.
       Two models are loaded.
       net_g_fix: pretrained model in source domain, used to guide the training of net_g
       net_g:
    """

    def init_training_settings(self):
        train_opt = self.opt['train']
        self.fix_backbone = self.opt['train'].get('fix_backbone',False)
        self.ema_decay = train_opt.get('ema_decay', 0)
        if self.ema_decay > 0:
            logger = get_root_logger()
            logger.info(f'Use Exponential Moving Average with decay: {self.ema_decay}')
            # define network net_g with Exponential Moving Average (EMA)
            # net_g_ema is used only for testing on one GPU and saving
            # There is no need to wrap with DistributedDataParallel
            self.net_g_ema = build_network(self.opt['network_g']).to(self.device)
            # load pretrained model
            load_path = self.opt['path'].get('pretrain_network_g', None)
            if load_path is not None:
                self.load_network(self.net_g_ema, load_path, self.opt['path'].get('strict_load_g', True), 'params_ema')
            else:
                self.model_ema(0)  # copy net_g weight
            self.net_g_ema.eval()
        
        ### use ema or net_g_fixed
        ## define network net_g_fix
        self.use_net_g_ema = train_opt.get('use_net_g_ema', False)
        if not self.use_net_g_ema:
            self.net_g_fix = build_network(self.opt['network_g_fix'])
            self.net_g_fix = self.model_to_device(self.net_g_fix)
            # load pretrained models
            load_path = self.opt['path'].get('pretrain_network_g_fix', None)
            if load_path is not None:
                param_key = self.opt['path'].get('param_key_g_fix', 'params')
                model_type = self.opt['network_g_fix']['type']
                self.load_network(self.net_g_fix, load_path, self.opt['path'].get('strict_load_g_fix', True), param_key, model_type)
        
        # define network net_d
        self.net_d = build_network(self.opt['network_d'])
        self.net_d = self.model_to_device(self.net_d)
        self.print_network(self.net_d)

        # load pretrained models
        load_path = self.opt['path'].get('pretrain_network_d', None)
        if load_path is not None:
            param_key = self.opt['path'].get('param_key_d', 'params')
            self.load_network(self.net_d, load_path, self.opt['path'].get('strict_load_d', True), param_key)

        self.net_g.train()
        self.net_d.train()
        
        ### use ema or net_g_fixed

        # define losses
        if train_opt.get('pixel_opt'):
            self.cri_pix = build_loss(train_opt['pixel_opt']).to(self.device)
        else:
            self.cri_pix = None

        if train_opt.get('perceptual_opt'):
            self.cri_perceptual = build_loss(train_opt['perceptual_opt']).to(self.device)
        else:
            self.cri_perceptual = None

        if train_opt.get('gan_opt'):
            self.cri_gan = build_loss(train_opt['gan_opt']).to(self.device)
        
        if train_opt.get('domain_opt'):
            self.cri_domain = build_loss(train_opt['domain_opt']).to(self.device)
        else:
            self.cri_domain = None

        self.net_d_iters = train_opt.get('net_d_iters', 1)
        self.net_d_init_iters = train_opt.get('net_d_init_iters', 0)
        self.domain_opt_warmup_iters = train_opt.get('domain_opt_warmup_iters', 0) + self.net_d_init_iters
        self.domain_opt_detach = train_opt.get('domain_opt_detach', True)
        self.domain_opt_wogan = train_opt.get('domain_opt_wogan', False)
        # set up optimizers and schedulers
        self.setup_optimizers()
        self.setup_schedulers()
    
    def load_network(self, net, load_path, strict=True, param_key='params', model_type=''):
        """Load network.

        Args:
            load_path (str): The path of networks to be loaded.
            net (nn.Module): Network.
            strict (bool): Whether strictly loaded.
            param_key (str): The parameter key of loaded network. If set to
                None, use the root 'path'.
                Default: 'params'.
        """
        if model_type == '':
            model_type = self.opt['network_g']['type']
        logger = get_root_logger()
        net = self.get_bare_model(net)
        load_net = torch.load(load_path, map_location=lambda storage, loc: storage)
        if param_key is not None:
            if param_key not in load_net and 'params' in load_net:
                param_key = 'params'
                logger.info('Loading: params_ema does not exist, use params.')
            elif param_key not in load_net and 'params' not in load_net:
                param_key = ''
            if param_key != '':
                load_net = load_net[param_key]
        logger.info(f'Loading {net.__class__.__name__} model from {load_path}, with param key: [{param_key}].')
        # remove unnecessary 'module.'
        for k, v in deepcopy(load_net).items():
            if k.startswith('module.'):
                load_net[k[7:]] = v
                load_net.pop(k)
                k=k[7:]
                
        self._print_different_keys_loading(net, load_net, strict)
        net.load_state_dict(load_net, strict=strict)
    
    def setup_optimizers(self):
        train_opt = self.opt['train']
        # optimizer g
        optim_type = train_opt['optim_g'].pop('type')
        self.optimizer_g = self.get_optimizer(optim_type, self.net_g.parameters(), **train_opt['optim_g'])
        self.optimizers.append(self.optimizer_g)
        # optimizer d
        optim_type = train_opt['optim_d'].pop('type')
        self.optimizer_d = self.get_optimizer(optim_type, self.net_d.parameters(), **train_opt['optim_d'])
        self.optimizers.append(self.optimizer_d)
    
    def optimize_parameters(self, current_iter):
        # optimize net_g
        for p in self.net_d.parameters():
            p.requires_grad = False

        self.optimizer_g.zero_grad()
        
        ## Get prediction from fixed network
        ## lq B, C, H, W
        if self.use_net_g_ema:
            with torch.no_grad():
                self.output_s_fix, self.output_t_fix = self.net_g_ema(torch.cat([self.lq_s, self.lq_t], dim=0)).chunk(2, dim=0)
        else:    
            with torch.no_grad():
                self.output_s_fix, self.output_t_fix = self.net_g_fix(torch.cat([self.lq_s, self.lq_t], dim=0)).chunk(2, dim=0)
        self.output_s, self.output_t = self.net_g(torch.cat([self.lq_s, self.lq_t], dim=0)).chunk(2, dim=0)

        l_g_total = 0
        loss_dict = OrderedDict()
        if (current_iter % self.net_d_iters == 0 and current_iter > self.net_d_init_iters):
            # pixel loss
            if self.cri_domain and (current_iter > self.domain_opt_warmup_iters):
                if self.domain_opt_detach:
                    l_g_inter, l_g_intra = self.cri_domain(self.output_t, self.output_t_fix.detach(), 
                                                self.output_s.detach(), self.output_s_fix.detach())
                else:
                    l_g_inter, l_g_intra = self.cri_domain(self.output_t, self.output_t_fix.detach(), 
                                                self.output_s, self.output_s_fix.detach())
                
                if l_g_inter is not None:
                    l_g_total += l_g_inter
                    loss_dict['l_inter'] = l_g_inter
                if l_g_intra is not None:
                    l_g_total += l_g_intra
                    loss_dict['l_intra'] = l_g_intra
            if self.cri_pix:
                l_g_pix = self.cri_pix(self.output_s, self.gt_s)
                l_g_total += l_g_pix
                loss_dict['l_g_pix'] = l_g_pix
                
            if isinstance(self.output_s_fix, tuple):
                self.output_s_fix,_ = self.output_s_fix
            # perceptual loss
            if self.cri_perceptual:
                l_g_percep, l_g_style = self.cri_perceptual(self.output_s, self.gt_s)
                if l_g_percep is not None:
                    l_g_total += l_g_percep
                    loss_dict['l_g_percep'] = l_g_percep
                if l_g_style is not None:
                    l_g_total += l_g_style
                    loss_dict['l_g_style'] = l_g_style
            # gan loss
            if self.domain_opt_wogan or (current_iter <= self.domain_opt_warmup_iters):
                fake_g_pred_t = self.net_d(self.output_s)
            else:
                fake_g_pred_t = self.net_d(torch.cat([self.output_s, self.output_t], dim=0))
            l_g_gan = self.cri_gan(fake_g_pred_t, True, is_disc=False)
            l_g_total += l_g_gan
            loss_dict['l_g_gan'] = l_g_gan

            l_g_total.backward()
            self.optimizer_g.step()

        # optimize net_d
        for p in self.net_d.parameters():
            p.requires_grad = True

        self.optimizer_d.zero_grad()
        # real
        real_d_pred = self.net_d(self.gt_s.detach())
        l_d_real = self.cri_gan(real_d_pred, True, is_disc=True)
        loss_dict['l_d_real'] = l_d_real
        loss_dict['out_d_real'] = torch.mean(real_d_pred.detach())
        l_d_real.backward()
        # fake
        fake_d_pred = self.net_d(self.output_s.detach())
        l_d_fake = self.cri_gan(fake_d_pred, False, is_disc=True)
        loss_dict['l_d_fake'] = l_d_fake
        loss_dict['out_d_fake'] = torch.mean(fake_d_pred.detach())
        l_d_fake.backward()
        self.optimizer_d.step()

        self.log_dict = self.reduce_loss_dict(loss_dict)

        if self.ema_decay > 0:
            self.model_ema(decay=self.ema_decay)

    def save(self, epoch, current_iter):
        if hasattr(self, 'net_g_ema'):
            self.save_network([self.net_g, self.net_g_ema], 'net_g', current_iter, param_key=['params', 'params_ema'])
        else:
            self.save_network(self.net_g, 'net_g', current_iter)
        self.save_network(self.net_d, 'net_d', current_iter)
        self.save_training_state(epoch, current_iter)
